import streamlit as st
import os
import tempfile
import shutil
from pathlib import Path
import logging
from image_style_trans.style_processor import StyleProcessor

class StyleUI:
    def __init__(self, project_root, output_dir):
        """
        初始化风格转换UI
        :param project_root: 项目根目录
        :param output_dir: 输出目录
        """
        self.project_root = Path(project_root)
        self.output_dir = Path(output_dir)
        self.style_processor = StyleProcessor()
        
        # 创建临时目录
        self.temp_dir = tempfile.mkdtemp()
        
        # 注册清理函数
        import atexit
        atexit.register(self.cleanup)
    
    def cleanup(self):
        """
        清理临时文件
        """
        try:
            shutil.rmtree(self.temp_dir)
        except Exception as e:
            logging.error(f"清理临时文件时发生错误: {str(e)}")
    
    def render(self):
        """
        渲染UI界面
        """
        st.header("图像风格转换")
        st.markdown("""
        使用可灵AI的图生图API对图像进行风格转换。
        上传一张图片，并输入风格描述，即可生成具有新风格的图像。
        """)
        
        # 检查服务可用性
        if not self.style_processor.available:
            st.error("图像风格转换服务未正确初始化，请检查API密钥和OSS配置")
            if not self.style_processor.kelin_service.available:
                st.info("请设置KELIN_API_KEY环境变量")
            if not self.style_processor.oss_manager.available:
                st.info("请设置OSS_ACCESS_KEY_ID和OSS_ACCESS_KEY_SECRET环境变量")
            return
        
        # 上传图像
        uploaded_file = st.file_uploader("上传图像", type=["jpg", "jpeg", "png"])
        
        # 风格设置
        prompt = st.text_area("风格描述", placeholder="例如：油画风格，梵高风格，水彩画风格等")
        negative_prompt = st.text_area("负向提示词（可选）", placeholder="例如：模糊，扭曲，低质量等")
        
        # 处理按钮
        process_button = st.button("开始处理")
        
        # 创建结果占位符
        result_placeholder = st.empty()
        
        if process_button and uploaded_file is not None and prompt:
            # 显示处理中提示
            with st.spinner("正在处理图像，请稍候...这可能需要一些时间"):
                # 保存上传的文件
                temp_file_path = os.path.join(self.temp_dir, uploaded_file.name)
                with open(temp_file_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                # 创建输出目录
                style_output_dir = self.output_dir / "styled_images"
                os.makedirs(style_output_dir, exist_ok=True)
                
                # 处理图像
                result_path = self.style_processor.process_image(
                    image_path=temp_file_path,
                    prompt=prompt,
                    negative_prompt=negative_prompt if negative_prompt else None,
                    output_dir=style_output_dir
                )
                
                if result_path:
                    # 显示结果
                    result_placeholder.success(f"图像处理成功！")
                    
                    # 创建两列布局
                    col1, col2 = st.columns(2)
                    
                    # 显示原始图像
                    with col1:
                        st.subheader("原始图像")
                        st.image(temp_file_path)
                    
                    # 显示处理后的图像
                    with col2:
                        st.subheader("处理后的图像")
                        st.image(result_path)
                    
                    # 提供下载链接
                    with open(result_path, "rb") as f:
                        st.download_button(
                            label="下载处理后的图像",
                            data=f.read(),
                            file_name=os.path.basename(result_path),
                            mime="image/jpeg"
                        )
                else:
                    result_placeholder.error("图像处理失败，请检查日志获取详细信息")
        elif process_button:
            if not uploaded_file:
                st.warning("请上传图像")
            if not prompt:
                st.warning("请输入风格描述")

# 如果直接运行此文件，则启动Streamlit应用
if __name__ == "__main__":
    import sys
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    
    # 设置项目根目录和输出目录
    project_root = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    output_dir = project_root / "output"
    
    # 创建并渲染UI
    style_ui = StyleUI(project_root, output_dir)
    style_ui.render()